package edu.umd.clip.lm.model.decoding;

import java.util.*;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class VariationalTaggerTrellisSlice {
	Trie trie;
	
	VariationalTaggerTrellisSlice() {
		trie = new Trie(-1);
	}
	
	static class Trie implements Comparable<Trie> {
		Trie subTries[];
		final int tag;
		double prob;
		
		public Trie(int tag, double prob) {
			this.tag = tag;
			this.prob = prob;
		}

		public Trie(int tag) {
			this(tag, 0);
		}
		
		public int compareTo(Trie other) {
			return tag - other.tag;
		}
		
		public Trie getSubtrie(int tag) {
			if (subTries == null) return null;
			for(int i=0; i<subTries.length; ++i) {
				if (subTries[i].tag == tag) {
					return subTries[i];
				}
			}
			return null;
		}
		
		public void addSubtrie(Trie t) {
			if (subTries == null) {
				subTries = new Trie[1];
			} else {
				subTries = Arrays.copyOf(subTries, subTries.length+1);
			}
			subTries[subTries.length-1] = t;
		}
		
		/**
		 * @param ctx the context is organized as follows: tag[i-1], t[i-2], ...
		 * 
		 * @return
		 */
		public Trie findTrie(int ctx[]) {
			return findTrie(ctx, 0);
		}
		
		Trie findTrie(int ctx[], int idx) {
			int t = ctx[idx];
			Trie subTrie = getSubtrie(t);
			if (idx == ctx.length - 1) {
				return subTrie;
			}
			if (subTrie == null) return null;
			return subTrie.findTrie(ctx, idx+1);
		}
		
		/**
		 * @param ctx the context is organized as follows: tag[i-1], t[i-2], ..., t[i]
		 * @param prob
		 */
		public void addProb(int ctx[], double prob) {
			addProb(ctx, 0, prob);
		}
		
		void addProb(int ctx[], int idx, double prob) {
			this.prob += prob;
			int tag = ctx[idx];
			
			// search for tag
			Trie subTrie = getSubtrie(tag);
			if (subTrie == null) {
				if (idx == ctx.length - 1) {
					// add item
					Trie t = new Trie(tag, prob);
					addSubtrie(t);
					return;
				}
				// add sub trie
				subTrie = new Trie(tag);
				addSubtrie(subTrie);
			} else if (idx == ctx.length - 1) {
				subTrie.prob += prob;
				return;
			}
			subTrie.addProb(ctx, idx+1, prob);
		}

		/**
		 * adds probability to the tag subtrie, creates the subtrie if necessary
		 * @param tag
		 * @param prob
		 * @return the tag subtrie
		 */
		public Trie addProb(int tag, double prob) {
			Trie subTrie = getSubtrie(tag);
			if (subTrie == null) {
				subTrie = new Trie(tag, prob);
				addSubtrie(subTrie);
				return subTrie;
			}
			subTrie.prob += prob;
			return subTrie;
		}
	}

	public void addProb(int[] ctx, double prob) {
		trie.addProb(ctx, prob);
	}

	public Trie findTrie(int[] ctx) {
		return trie.findTrie(ctx);
	}
	
}
